import pandas as pd
from . import config, utils
import logging
from sklearn.model_selection import KFold
import pickle


def main():
    assert config.dataset == 'val'
    train = pd.read_csv(config.pj_root + 'data/' + config.dataset + '.csv', index_col='no')

    assert train.columns[0] == 'flag'
    a_feature = pd.read_csv(config.pj_root + 'data/a_feature.csv', index_col='no')
    train = train.join(a_feature)

    if len(config.drop_columns) != 0:
        logging.warning('Will drop %s len=%d' % (str(config.drop_columns), len(config.drop_columns)))
        train = train.drop(config.drop_columns, axis=1)

    if len(config.select_columns) != 0:
        config.select_columns.insert(0, 'flag')  # use flag in training
        logging.warning('Will select %s len=%d' % (str(config.select_columns), len(config.select_columns)))
        train = train[train.columns[train.columns.isin(config.select_columns)]]
    else:
        config.select_columns = list(train.columns)[1:]



    if config.use_basic_process:
        train, col_func_map = utils.basic_process(train, has_flag=True)
        utils.dump_to_data(col_func_map, 'col_func_map.pkl')

    X = train.values[:, 1:]
    Y = train.values[:, 0:1].reshape((-1,))

    if config.task == 'val':
        logging.info('Use %d Folds...' % config.kfold_k)
        logging.info(config.model_para)
        kf = KFold(n_splits=config.kfold_k)
        all_score = 0
        for i, (train_index, test_index) in enumerate(kf.split(X)):
            X_train, X_test = X[train_index], X[test_index]
            y_train, y_test = Y[train_index], Y[test_index]
            model = config.model(**config.model_para)
            model.fit(X_train, y_train)
            y_pred = model.predict_proba(X_test)
            y_pred = y_pred[:, 1]
            score = utils.report(y_test, y_pred)
            all_score += score
            logging.info('Fold %d/%d Score: %f ' % (i + 1, config.kfold_k, score))
            # i = pd.DataFrame({'feature': train.columns[1:], 'importance': model.feature_importances_})
            # i = i[i.importance > 0].sort_values('importance', ascending=False)
            # print(i)
            # print(len(i))
        logging.info('Avg score %f. ' % (all_score / config.kfold_k))
    elif config.task == 'valall':
        X_train = X
        y_train = Y
        logging.info('Start training....')
        model = config.model(**config.model_para)
        model.fit(X_train, y_train)
        with open(config.pj_root + 'model/%s.mo' % config.model.__name__, 'wb') as f:
            pickle.dump(model, f)
        # i = pd.DataFrame({'feature': train.columns[1:], 'importance': model.feature_importances_})
        # i = i[i.importance > 0].sort_values('importance', ascending=False)
        # print(i)
        # print(len(i))
    else:
        raise Exception('Unknown Task~!!')

if __name__ == '__main__':
    utils.could_override_config()  # default override model_para
    main()
